import os
import gradio as gr
from openai import OpenAI

# os.environ["HF_HOME"] = "~/.cache/gitee-ai"
# os.environ["HF_ENDPOINT"] = "https://hf-api.gitee.com"

# from transformers import AutoTokenizer, AutoModel
# tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
# model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
# model.eval()

# 读取环境私钥
G_KEY = os.getenv("CHATGLM_KEY")

print(G_KEY)

# 配置密钥和请求地址
client = OpenAI(
    api_key=G_KEY,
    base_url="https://open.bigmodel.cn/api/paas/v4/",
    # base_url="https://api.moonshot.cn/v1",
)


def chat_with_gpt(prompt):
    print("对话开始", prompt)
    response = client.chat.completions.create(
        model="glm-4-flash",  # 根据当前上下文占用的 Tokens 数量来选择合适的模型 8k 32k 128k
        messages=[{"role": "user", "content": prompt}],
        # 使用什么采样温度，介于 0 和 1 之间。较高的值（如 0.7）将使输出更加随机，而较低的值（如 0.2）将使其更加集中和确定性
        temperature=0.95,
        # 是否流式返回
        # stream=True,
        # tool_calls 调用
        # tools=tools,
    )
    print("对话结果：")
    print(response)
    return response.choices[0].message.content


# Gradio 界面
iface = gr.Interface(
    fn=chat_with_gpt,
    inputs="text",
    outputs="text",
    title="法小猿Chatbot",
    description="法小猿提供法律咨询服务，立即进行对话",
    examples=[
        "中国民法典规定自然人有哪些权利",
        "买卖人体细胞属于犯罪吗",
        "侵犯他人肖像权将受到什么惩罚",
    ],
)


if __name__ == "__main__":
    iface.launch()
